🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125850

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 3 Unrelated Failures

As of commit a46722a with merge base 8f30f36 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

/cc @ezyang @eqy This is an explorative blackbox PR as I don't have free cuda resources right now and we don't have a quick way to setup the env to contribute sparse c++/cuda PR (see #125297).

But I made this editable on your side in the case you have the env ready and a quick fix is enough.

The committers listed above are authorized under a signed CLA.

• ✅ login: bhack (0841181, bb7240c, a47573b, 637532c, 65dbb6c, 8cce2c7, 79ad3d3, 052bb61, a2554de, 679b9fd, dc29c0f, 5332381, a46722a, 2362079, 58eeb41, faa531f)

Should we add a (presumably large tensor) test for this?

Do we had an INT_MAX test already somewhere that we could expand?

Unfortunately these are not really unified at the moment, but this should surface some examples: https://github.com/search?q=repo%3Apytorch%2Fpytorch+64bit+language%3APython+path%3A%2F%5Etest%5C%2F%2F&type=code

As we don't have a specific CUDA test do we want to find a workaround from python?

Can you suggest one from grep -R torch.nonzero test/?

I think I am going to close this as cub::DispatchSelectIf probably it will be slower then cub::DeviceSelect::Flagged we are currently using.

Probably we need to wait upstream for NVIDIA/cccl#1422

What do you think?

@ezyang Do you think we can we open a new ticket to lower this with Trition where and sum?
https://github.com/pytorch/pytorch/blob/a174c536f8f32b41da9efa647e364196423468a5/torch/_inductor/lowering.py#L2187C20-L2187C35

Edit:
The ticket is at #126003

yes, need a big tensor test. @eqy's link is good for examples

Ok thanks,
so I am going to close it as I don't have the env and currently spare GPU computing to write a brand new test and recompile it.
At least if we don't identify another python test that it is already indirectly using nonzero and that we could modify it with a big input.

Just to check if it could compile at least with the current CUB version.
Do you know what is this CI failure?

/usr/local/cuda/include/cub/agent/agent_select_if.cuh(264): error: function "at::native::<unnamed>::NonZeroOp<T>::operator() [with T=c10::complex<c10::Half>]" cannot be called with the given argument list
            argument types are: (int64_t)
            object type is: at::native::<unnamed>::NonZeroOp<c10::complex<c10::Half>>
                  selection_flags[ITEM] = select_op(items[ITEM]);

I think we need cub/cub/agent/agent_select_if.cuh changes introduced at NVIDIA/cccl#1379

So this mean that we need to wait for the next cuda 12.4 update and make it also conditional.

This PR seems fine. I agree you may need to preprocessor your way to victory. CI will say.

@ezyang the new CUDA 12.5 delivers CUB 2.4.0 so it could be enough for this workaround.

if you're willing to wait for cuda 12.5 :)

if you're willing to wait for cuda 12.5 :)

This version is required for the workaround API. A full upstream solution it will require to wait more CUDA releases.

Sorry, I'm not sure I understand the state of this PR. I would happily accept a PR that makes nonzero int64 work on CUDA 12.5 or later, subject to the requirement that this functionality is preprocessored out. If you don't mind waiting, we can also ice this PR until CUDA 12.6 shows up by default and then we can just land it as is.

The current status for testing/using the workaround in this PR is to have CUB 2.4.0.
As in pytorch we currently use CUB from the official CUDA distribution it means we need upgrade pytorch to support CUDA 12.5 (any plan?).

For a "regular" upstream solution we need that this will be merged:
NVIDIA/cccl#1422

And after that we need an official CUB release and that release will be included in a CUDA release.
So we have no guaranties that we will have an upstream solution with CUDA 12.6

@ezyang See NVIDIA/cccl#1454 (comment)

@atalman for CUDA 12.5